跳转至

用树结构打破链条：具有 $mathcal{O}(log N)$ 时间复杂度的模型并行深度学习

文章背景与核心概要

传统的深度神经网络训练严重依赖误差反向传播（Backpropagation），这要求网络层与层之间进行严格串行的前向计算和后向传播。随着网络深度的不断增加，这种层间强依赖性构成了主要的性能瓶颈，导致训练过程只能局限于数据并行方案，而无法实现真正的高效模型并行。

为了突破这一限制，本文引入了 TreeProp——一种与具体架构无关的变分学习框架。该框架创新性地将网络层组织为树状层级结构，用分层计算取代了传统的串行前向与后向梯度传播，从而将一个 \(N\) 层网络的并行时间复杂度大幅降至 \(\mathcal{O}(\log N)\)。据我们所知，这是首个在训练的前向计算和后向梯度传播中均实现对数级并行时间复杂度的深度神经网络学习算法。

除了突破训练的并行性瓶颈外，TreeProp 还能在无需额外训练开销的前提下，隐式学习到具有不同有效深度的子网络。实验表明，该方法在视觉分类和自回归语言建模任务中均能达到与传统端到端训练相媲美的性能，并优于以往的对比学习方法，同时还能自然地推广至依赖沿时间反向传播（BPTT）的循环神经网络中。

论文元数据 (Paper Metadata)

  • arXiv ID: 2606.21497v2 [cs.LG]
  • 学科分类 (Subjects): 机器学习 (cs.LG)；人工智能 (cs.AI)；数据结构与算法 (cs.DS)
  • 作者 (Authors):
  • Neeraj Mohan Sushma
  • Aditya Nagarsekar
  • Cabrel Teguemne Fokam
  • Robin Schiewer
  • Amit Kumar Pal
  • Anand Subramoney
  • David Kappel
  • 提交日期 (Submission Date): 2026年6月19日（最后修订于：2026年8月14日）
  • 许可协议 (License): 知识共享署名 4.0 国际许可协议 license icon

摘要 (Abstract)

现代深度神经网络通常使用误差反向传播进行训练，这需要跨网络层进行串行的前向和后向计算。随着这些网络变得越来越深，这种方式带来了局限性，因为逐层更新是严格相互依存的，无法并行进行。这些约束将训练过程限制在数据并行方案中，从而排除了模型平行的可能性。

我们提出了 TreeProp，这是一个与架构无关的变分学习框架，它将网络层组织成树状的层级结构。在训练过程中，TreeProp 用分层计算取代了串行的前向计算和后向梯度传播。对于一个具有 \(N\) 层的网络，这使得中间表示和学习信号能够在 \(\mathcal{O}(\log N)\) 的时间复杂度内构建完成。据我们所知，TreeProp 是第一个在训练过程中的前向计算和后向梯度传播都具有对数并行时间复杂度的深度神经网络学习算法。

此外，我们表明该层级结构中存在多条有效的路径，使得 TreeProp 能够隐式学习具有不同有效深度的子网络，且无需额外的训练成本。我们在视觉分类和自回归语言建模上对 TreeProp 进行了评估，它在各种任务中均匹配了传统端到端训练的性能，并优展了以往的对比学习方法。我们进一步展示了 TreeProp 在循环神经网络（通常依赖于沿时间反向传播）中的适用性。

Modern deep neural networks are trained using error backpropagation, which requires sequential forward and backward computations across network layers. As these networks become deeper, this introduces limitations, since layer-wise updates are strictly interdependent and cannot proceed in parallel. These constraints restrict training procedures to data-parallel schemes, thereby prohibiting model-parallel training.

We propose TreeProp, an architecture-agnostic variational learning framework that organizes network layers into a tree-structured hierarchy. During training, TreeProp replaces sequential forward computations and backward gradient propagation with hierarchical computations. This allows intermediate representations and learning signals to be constructed in time complexity of \(\mathcal{O}(\log N)\) for a network of \(N\) layers. To the best of our knowledge, TreeProp is the first learning algorithm for deep neural networks with logarithmic parallel time complexity for both forward computation and backward gradient propagation during training.

Furthermore, we show that multiple valid paths through the hierarchy exist, such that TreeProp implicitly learns subnetworks with different effective depths, but without additional training effort. We evaluate TreeProp on vision classification and autoregressive language modeling, matching the performance of conventional end-to-end training for a variety of tasks and outperforming previous contrastive training approaches. We further demonstrate the applicability of TreeProp to recurrent neural networks that otherwise rely on backpropagation through time.